"""
数据工具函数
包含数据加载、特征提取等功能
"""
import os
import ast
import numpy as np
import pandas as pd
from PIL import Image
from typing import List, Dict, Any


def parse_list_string(s):
    """
    解析字符串形式的列表
    
    参数:
        s: 字符串，如 "[1, 2, 3]"
    
    返回:
        list: 解析后的列表
    
    示例:
        >>> parse_list_string("[1, 2, 3]")
        [1, 2, 3]
    """
    try:
        return ast.literal_eval(s)
    except:
        return []


def extract_meniscus_features(img_array, pixel_scale):
    """
    从弯液面图像中提取传统特征
    
    参数:
        img_array: 图像数组 (H, W)，值范围 [0, 1]
        pixel_scale: 像素尺度
    
    返回:
        features: 包含6个特征的列表
            [area, sum_gray, mean_gray, entropy, aspect_ratio, contrast]
    """
    try:
        # 提取非零区域（弯液面区域）
        meniscus_region = img_array[img_array > 0]
        
        # 1. 面积特征
        area = len(meniscus_region) * pixel_scale * pixel_scale
        
        # 2. 灰度和
        sum_gray = np.sum(meniscus_region) * pixel_scale * pixel_scale
        
        # 3. 平均灰度
        mean_gray = np.mean(meniscus_region) if len(meniscus_region) > 0 else 0
        
        # 4. 熵特征
        hist, _ = np.histogram(img_array.flatten(), bins=64, range=(0, 1))
        hist = hist[hist > 0]
        hist = hist / (hist.sum() + 1e-10)
        entropy = -np.sum(hist * np.log2(hist + 1e-10))
        
        # 5. 宽高比
        binary = img_array > 0
        indices = np.where(binary)
        if len(indices[0]) > 0:
            height = indices[0].max() - indices[0].min() + 1
            width = indices[1].max() - indices[1].min() + 1
            aspect_ratio = width / max(height, 1)
        else:
            aspect_ratio = 1
        
        # 6. 对比度（标准差）
        contrast = np.std(meniscus_region) if len(meniscus_region) > 0 else 0
        
        return [area, sum_gray, mean_gray, entropy, aspect_ratio, contrast]
    
    except Exception as e:
        print(f"Warning: Feature extraction error - {e}")
        return [0, 0, 0, 0, 1, 0]


def load_csv_data(data_dir, csv_files):
    """
    从CSV文件加载数据
    
    参数:
        data_dir: 数据目录
        csv_files: CSV文件名列表
    
    返回:
        all_data: 数据字典列表，每个元素包含:
            - image_paths: 图像路径列表
            - pixel_scales: 像素尺度列表
            - meniscus_volume: 弯液面体积
    """
    all_data = []
    
    for csv_file in csv_files:
        file_path = os.path.join(data_dir, csv_file)
        
        if not os.path.exists(file_path):
            print(f"Warning: File not found - {file_path}")
            continue
        
        print(f"Loading: {csv_file}")
        
        try:
            df = pd.read_csv(file_path, encoding='utf-8')
        except:
            try:
                df = pd.read_csv(file_path, encoding='gbk')
            except:
                print(f"  Error: Cannot read file with utf-8 or gbk encoding")
                continue
        
        valid_count = 0
        invalid_count = 0
        
        for idx, row in df.iterrows():
            try:
                image_paths = parse_list_string(row['ImagePaths'])
                pixel_scales = parse_list_string(row['PixelScale'])
                
                # 验证数据完整性
                if len(image_paths) != 20 or len(pixel_scales) != 20:
                    invalid_count += 1
                    continue
                
                data_item = {
                    'image_paths': image_paths,
                    'pixel_scales': pixel_scales,
                    'meniscus_volume': float(row['MeniscusVolume'])
                }
                all_data.append(data_item)
                valid_count += 1
                
            except Exception as e:
                invalid_count += 1
                continue
        
        print(f"  Valid: {valid_count}, Invalid: {invalid_count}, Total: {len(df)}")
    
    print(f"\nTotal loaded: {len(all_data)} samples")
    return all_data


def load_and_preprocess_image(image_path, pixel_scale, image_size=(288, 96)):
    """
    加载并预处理单张图像
    
    参数:
        image_path: 图像路径
        pixel_scale: 像素尺度
        image_size: 目标图像尺寸 (width, height)
    
    返回:
        img_array: 预处理后的图像数组 [H, W]
        features: 提取的特征列表 [6]
    """
    try:
        if os.path.exists(image_path):
            img = Image.open(image_path).convert('L')
            img = img.resize(image_size, Image.BILINEAR)
            img_array = np.array(img).astype(np.float32) / 255.0
        else:
            # 图像不存在时使用零数组
            img_array = np.zeros(image_size[::-1], dtype=np.float32)
    except Exception as e:
        print(f"Error loading image {image_path}: {e}")
        img_array = np.zeros(image_size[::-1], dtype=np.float32)
    
    # 提取特征
    features = extract_meniscus_features(img_array, pixel_scale)
    
    return img_array, features


def compute_dataset_statistics(data_list):
    """
    计算数据集统计信息
    
    参数:
        data_list: 数据列表
    
    返回:
        stats: 统计信息字典
    """
    volumes = [item['meniscus_volume'] for item in data_list]
    
    stats = {
        'num_samples': len(volumes),
        'volume_min': np.min(volumes),
        'volume_max': np.max(volumes),
        'volume_mean': np.mean(volumes),
        'volume_std': np.std(volumes),
        'volume_median': np.median(volumes),
    }
    
    return stats


def print_dataset_statistics(train_data, val_data):
    """
    打印数据集统计信息
    
    参数:
        train_data: 训练数据
        val_data: 验证数据
    """
    print("\n" + "="*70)
    print("Dataset Statistics".center(70))
    print("="*70)
    
    train_stats = compute_dataset_statistics(train_data)
    val_stats = compute_dataset_statistics(val_data)
    
    print("\n[Training Set]")
    for key, value in train_stats.items():
        if 'num' in key:
            print(f"  {key}: {value}")
        else:
            print(f"  {key}: {value:.4f}")
    
    print("\n[Validation Set]")
    for key, value in val_stats.items():
        if 'num' in key:
            print(f"  {key}: {value}")
        else:
            print(f"  {key}: {value:.4f}")
    
    print("="*70 + "\n")


if __name__ == "__main__":
    # 测试功能
    print("Testing data utils...")
    
    # 测试特征提取
    print("\n1. Testing feature extraction...")
    test_img = np.random.rand(96, 288).astype(np.float32)
    features = extract_meniscus_features(test_img, 0.1)
    print(f"   Extracted features: {len(features)}")
    print(f"   Features: {[f'{f:.4f}' for f in features]}")
    
    # 测试列表解析
    print("\n2. Testing list parsing...")
    test_str = "[1.0, 2.0, 3.0]"
    parsed = parse_list_string(test_str)
    print(f"   Input: {test_str}")
    print(f"   Parsed: {parsed}")
    
    # 测试统计
    print("\n3. Testing statistics...")
    test_data = [
        {'meniscus_volume': 10.0},
        {'meniscus_volume': 15.0},
        {'meniscus_volume': 20.0},
    ]
    stats = compute_dataset_statistics(test_data)
    print(f"   Stats: {stats}")
    
    print("\n✓ Utils test passed!")